#ifndef YOLO5_PRE_H
#define YOLO5_PRE_H

#include <opencv2/opencv.hpp>
#include <string>
#include <vector>
#include <memory>
#include "rga2_wrapper.h"

// 固定尺寸常量定义
namespace Yolo5PreConstants {
    const int INPUT_WIDTH = 480;      // 输入图像宽度
    const int INPUT_HEIGHT = 480;     // 输入图像高度
    const int OUTPUT_SIZE = 640;      // 输出图像尺寸
    const size_t NV12_Y_SIZE = INPUT_WIDTH * INPUT_HEIGHT;           // Y分量大小: 230,400字节
    const size_t NV12_UV_SIZE = NV12_Y_SIZE / 2;                    // UV分量大小: 115,200字节
    const size_t NV12_TOTAL_SIZE = NV12_Y_SIZE + NV12_UV_SIZE;      // 总大小: 345,600字节
}

class Yolo5Pre {
public:
    Yolo5Pre();
    ~Yolo5Pre();
    
    // 处理数据集目录
    bool processDataset(const std::string& sourceDir, const std::string& destDir, double rotationAngle = 0.0);
    
    // 转换NV12格式的.raw文件为640x640的PNG
    bool convertNV12To640x640PNG(const std::string& rawFilePath, const std::string& outputDir, double rotationAngle = 0.0);
    
    // 将480x480的PNG图片填充到640x640
    bool padImageTo640x640(const std::string& imagePath, const std::string& destDir, double rotationAngle = 0.0);
    
    // 验证目录
    static bool validateDirectory(const std::string& directoryPath);

private:
    std::unique_ptr<RGA2Wrapper> m_rga2Wrapper;
    
    // 从NV12数据创建OpenCV Mat
    cv::Mat createMatFromNV12(const std::vector<uint8_t>& rawData, int width, int height);
    
    // 生成输出文件名
    std::string generateOutputFilename(const std::string& inputPath, const std::string& outputDir, 
                                     double rotationAngle, const std::string& suffix = "");
    
    // 创建输出目录
    bool createOutputDirectory(const std::string& outputDir);
    
    // 获取文件列表
    std::vector<std::string> getFileList(const std::string& directory, const std::string& extension);
};

#endif // YOLO5_PRE_H
